#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/11/5 14:35
# @USER    : Shengji He
# @File    : convert_model.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import glob
import os
import torch


def convert_model(src, dst):
    state_dict = torch.load(src, map_location=torch.device('cpu'))
    torch.save(state_dict, dst, _use_new_zipfile_serialization=False)


def main():
    glob_research = '../weights/Dataset291_TotalSegmentator_part1_organs_1559subj/nnUNetTrainerNoMirroring__nnUNetPlans__3d_fullres/fold_0/*.pth'
    files = glob.glob(glob_research, recursive=True)
    for file in files:
        dst = file.replace('.pth', '_lv.pth')
        if os.path.exists(dst):
            continue
        convert_model(file, dst)


if __name__ == '__main__':
    main()
    print('done')
